from construct2 import *


class IpAddressAdapter(Adapter):
    def _encode(self, obj):
        return "".join(chr(int(b)) for b in obj.split("."))
    def _decode(self, obj):
        return ".".join(str(ord(b)) for b in obj)

def IpAddress(name):
    return IpAddressAdapter(Bytes(name, 4))

ipv4_header = Struct("ip_header",
    EmbeddedBitStruct(
        OneOfValidator(
            Nibble("version"), 
            (4,)
        ),
        ExprAdapter(Nibble("header_length"), 
            decoder = lambda obj: obj * 4, 
            encoder = lambda obj: obj / 4
        ),
    ),
    BitStruct("tos",
        Bits("precedence", 3),
        Flag("minimize_delay"),
        Flag("high_throuput"),
        Flag("high_reliability"),
        Flag("minimize_cost"),
        Padding(1),
    ),
    UBInt16("total_length"),
    Value("payload_length", lambda ctx: ctx["total_length"] - ctx["header_length"]),
    UBInt16("identification"),
    EmbeddedBitStruct(
        Struct("flags",
            Padding(1),
            Flag("dont_fragment"),
            Flag("more_fragments"),
        ),
        Bits("frame_offset", 13),
    ),
    UBInt8("ttl"),
    Enum(UBInt8("protocol"),
        ICMP = 1,
        TCP = 6,
        UDP = 17,
    ),
    UBInt16("checksum"),
    IpAddress("source"),
    IpAddress("destination"),
    MetaField("options", lambda ctx: ctx["header_length"] - 20),
)


if __name__ == "__main__":
    cap = "4500003ca0e3000080116185c0a80205d474a126".decode("hex")
    obj = ip_header.parse(cap)
    print obj
    print repr(ip_header.build(obj))









